package tools

import (
	"context"
	"fmt"

	mcpgo "github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
	"github.com/safedep/dry/api/pb"
	"github.com/safedep/vet/mcp"
	"github.com/safedep/vet/pkg/common/logger"
)

// PackageMalwareTool provides functionality to scan packages for potential malware
type packageMalwareTool struct {
	driver mcp.Driver
}

// NewPackageMalwareTool creates a new instance of PackageMalwareTool
func NewPackageMalwareTool(driver mcp.Driver) *packageMalwareTool {
	return &packageMalwareTool{
		driver: driver,
	}
}

// Register registers the tool with the MCP server
func (t *packageMalwareTool) Register(server *server.MCPServer) error {
	tool := mcpgo.NewTool("get_package_version_malware_report",
		mcpgo.WithDescription("Get a malware analysis report for a package version"),
		mcpgo.WithString("purl", mcpgo.Required(), mcpgo.Description("The package URL to check for malware")),
	)

	server.AddTool(tool, t.executeCheckPackageMalware)
	return nil
}

func (t *packageMalwareTool) executeCheckPackageMalware(ctx context.Context,
	req mcpgo.CallToolRequest,
) (*mcpgo.CallToolResult, error) {
	purl, err := req.RequireString("purl")
	if err != nil {
		return nil, fmt.Errorf("purl is required: %w", err)
	}

	logger.Debugf("Checking package for malware: %s", purl)

	parsedPurl, err := pb.NewPurlPackageVersion(purl)
	if err != nil {
		return nil, fmt.Errorf("invalid purl: %w", err)
	}

	malwareReport, err := t.driver.GetPackageVersionMalwareReport(ctx, parsedPurl.PackageVersion())
	if err != nil {
		// Handle this error gracefully otherwise LLMs will hallucinate
		if err == mcp.ErrMaliciousPackageScanningPackageNotFound {
			return mcpgo.NewToolResultText("No known malware found for this package"), nil
		}

		return nil, fmt.Errorf("failed to get package version malware report: %w", err)
	}

	malwareReportJson, err := serializeForLlm(malwareReport)
	if err != nil {
		return nil, fmt.Errorf("failed to serialize malware report: %w", err)
	}

	return mcpgo.NewToolResultText(malwareReportJson), nil
}
